#include <stdio.h>
#include <cuda.h>
#include <sys/time.h>
#include <cuda_runtime.h>
#include <stdint.h>
#include "cublas_v2.h"

#define M 2048
#define N 2048
#define K 2048
#define TILE_SIZE 16
#define COMP 6

# define timersub(a, b, result)                                         \
        do {                                                            \
                (result)->tv_sec = (a)->tv_sec - (b)->tv_sec;           \
                (result)->tv_usec = (a)->tv_usec - (b)->tv_usec;        \
                if ((result)->tv_usec < 0) {                            \
                        --(result)->tv_sec;                             \
                        (result)->tv_usec += 1000000;                   \
                }                                                       \
        } while (0)

__global__ void kernel4(float *C, const float *A, const float *B, int m, int n,  int k)
{
	const int col = blockIdx.x * 96 + threadIdx.x;
	const int row = blockIdx.y * 96 + threadIdx.y;
	
	__shared__ float Ab[16 * COMP][16];
	__shared__ float Bb[16][16 * COMP];
	
	float sum[COMP * COMP];
	
	for (int i=0;i<k;i+=16) {
#pragma unroll
		for (int j=0;j<COMP;++j) {
			if ((i + threadIdx.x) >= k || (row + j * 16) >= m) {
				Ab[threadIdx.y + j * 16][threadIdx.x] = 0;
			}
			else {
				Ab[threadIdx.y + j * 16][threadIdx.x] = A[(row + j * 16) * k + i + threadIdx.x];
			}
		}
		
#pragma unroll
		for (int j=0;j<COMP;++j) {
			if ((i + threadIdx.y) >= k || (col + j * 16) >= n) {
				Bb[threadIdx.y][threadIdx.x + j * 16] = 0;
			}
			else {
				Bb[threadIdx.y][threadIdx.x + j * 16] = B[(i + threadIdx.y) * n + col + j * 16];
			}
		}
		
		__syncthreads();
		
#pragma unroll
		for (int x=0;x<16;++x) {
#pragma unroll
			for (int r=0;r<COMP;++r) {
#pragma unroll
				for (int c=0;c<COMP;++c) {
					sum[r * COMP + c] += Ab[threadIdx.y + r * 16][x] * Bb[x][threadIdx.x + c * 16];
				}
			}
		}
		__syncthreads();
	};

#pragma unroll
	for (int r=0;r<COMP;++r) {
#pragma unroll
		for (int c=0;c<COMP;++c) {
			if ( ((row + r * 16) < m) && ((col + c * 16)) < n ) {
				C[(row + r * 16) * n + col + c * 16] = sum[r * COMP + c];
			}
		}
	}
}

int main(int argc, const char *argv[])
{
	struct timeval time_start, time_end, time_res;

	float *m1_h, *m2_h, *o_h ;
	float *o_h2;
	float *m1_d, *m2_d, *o_d;

	float alpha = 1.0;
	float beta = 0.0;

	int count = 0;

        //int n = atoi(argv[1]);

	m1_h = (float *) malloc(M * K * sizeof(float));
	m2_h = (float *) malloc(K * N * sizeof(float));
	o_h = (float *) malloc(M * N * sizeof(float));
	o_h2 = (float *) malloc(M * N * sizeof(float));

	cudaMalloc( (void **)&m1_d, M * K * sizeof(float));
	cudaMalloc( (void **)&m2_d, K * N * sizeof(float));
	cudaMalloc( (void **)&o_d, M * N * sizeof(float));

	// Initialize host array and copy it to CUDA device
	for (int r = 0 ; r < M; ++r) {
		//o_h[x] = 0.0;
		for (int c = 0 ; c < K; ++c) {
			m1_h[r*K + c] = rand()/(float)RAND_MAX;
		}
	}

	for (int r = 0 ; r < K; ++r) {
		//o_h[x] = 0.0;
		for (int c = 0 ; c < N; ++c) {
			m2_h[r*N + c] = rand()/(float)RAND_MAX;
		}
	}

	cudaMemcpy(m1_d, m1_h, M * K * sizeof(float), cudaMemcpyHostToDevice);
	cudaMemcpy(m2_d, m2_h, K * N * sizeof(float), cudaMemcpyHostToDevice);

	cublasHandle_t handle;

	if(cublasCreate(&handle) != CUBLAS_STATUS_SUCCESS){
		printf("cublasCteate fail\n");
		exit(-1);
	}

	int i = TILE_SIZE;

	dim3 grid((M + 95)/96, (N + 95)/96, 1);
	dim3 threads(i, i, 1);

	gettimeofday(&time_start, NULL);
	
	kernel4 <<< grid, threads >>> ( o_d, m1_d, m2_d, M, N, K );

	cudaDeviceSynchronize();

	gettimeofday(&time_end, NULL);
	
	timersub(&time_end, &time_start, &time_res);
	
	printf("%ld,%03ld ms\n", time_res.tv_sec, time_res.tv_usec/1000);
	
	cudaMemcpy(o_h,o_d,M*N*sizeof(float), cudaMemcpyDeviceToHost);

	cudaDeviceSynchronize();

	gettimeofday(&time_start,NULL);

	cublasSgemm(handle, CUBLAS_OP_T, CUBLAS_OP_T, M, K, N, &alpha, m1_d, K, m2_d, N, &beta, o_d, N);
	cudaDeviceSynchronize();

	gettimeofday(&time_end,NULL);

	timersub(&time_end, &time_start, &time_res);

	printf("%ld,%03ld ms\n", time_res.tv_sec, time_res.tv_usec/1000);

	cudaMemcpy(o_h2,o_d, M*N*sizeof(float), cudaMemcpyDeviceToHost);

	for(int i = 0 ; i < M ; i ++){
		for(int j = 0 ; j < N ; j++){
			if( (o_h[i*N+j] - o_h2[j*N+i]) > 0.001 || (o_h[i*N+j] - o_h2[j*N+i]) < -0.001){
				count++;
			}
		}
	}
	printf("number of differences : %d\n",count);

	free(m1_h);
	free(m2_h);
	free(o_h);
	
	cudaFree(m1_d);
	cudaFree(m2_d);
	cudaFree(o_d);
}

